# 使用格式
# python object_tracker.py --prototxt deploy.prototxt --model res10_300x300_ssd_iter_140000.caffemodel

# import the necessary packages
from pyimagesearch.centroidtracker import CentroidTracker
from imutils.video import VideoStream
import numpy as np
import argparse
import imutils
import time
import cv2

# 构造参数解析对象
ap = argparse.ArgumentParser()
ap.add_argument("-p", "--prototxt", required=True,
	help="path to Caffe 'deploy' prototxt file")
ap.add_argument("-m", "--model", required=True,
	help="path to Caffe pre-trained model")
ap.add_argument("-c", "--confidence", type=float, default=0.5,
	help="minimum probability to filter weak detections")
args = vars(ap.parse_args())

# 初始化质心跟踪器和视频帧的尺寸
ct = CentroidTracker()
(H, W) = (None, None)

# 从磁盘加载序列化模型
print("[INFO] loading model...")
net = cv2.dnn.readNetFromCaffe(args["prototxt"], args["model"])

# 初始化视频流并允许传感器预热
print("[INFO] starting video stream...")
vs = VideoStream(src=0).start()
time.sleep(2.0)

# 从视频流中循环提取视频帧
while True:
	# 从视频流中读取视频帧并调制视频帧尺寸大小
	frame = vs.read()
	frame = imutils.resize(frame, width=400)

	# 如果(H, W)为0，则给(H,W)f赋值
	if W is None or H is None:
		(H, W) = frame.shape[:2]

	# 在 frame 中构造一个 blob, 用于网络的前向传播。
	# 获得预测结果
	blob = cv2.dnn.blobFromImage(frame, 1.0, (W, H), (104.0, 177.0, 123.0)) # 构建blob
	net.setInput(blob)	# 传入网络
	detections = net.forward() # 进行输出预测
	print("detections\n", detections)
	
	# 循环遍历预测结果
	rects = []	# 初始化边界框矩形列表
	for i in range(0, detections.shape[2]):
		# 过滤掉低置信度结果
		if detections[0, 0, i, 2] > args["confidence"]:
			# 计算边界框的(x, y)像素坐标
			box = detections[0, 0, i, 3:7] * np.array([W, H, W, H])
			rects.append(box.astype("int"))

			# 可视化
			(startX, startY, endX, endY) = box.astype("int")
			cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2)

	# 使用rects更新质心跟踪器
	objects = ct.update(rects)

	# 循环遍历跟踪对象
	for (objectID, centroid) in objects.items():
		# draw both the ID of the object and the centroid of the
		# object on the output frame
		# 在输出帧中绘制对象id和质心
		text = "ID {}".format(objectID)
		cv2.putText(frame, text, (centroid[0] - 10, centroid[1] - 10), 
			cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
		cv2.circle(frame, (centroid[0], centroid[1]), 4, (0, 255, 0), -1)

	cv2.imshow("Frame", frame)
	key = cv2.waitKey(1) & 0xFF

	if key == ord("q"):
		break

# do a bit of cleanup
cv2.destroyAllWindows()
vs.stop()